package com.cn.algorithm.one.tree;

import java.util.Comparator;

public class RBTree<E> extends BBST<E> {

    public static final boolean RED = false;

    public static final boolean BLACK = true;

    public RBTree(Comparator<E> comparator) {
        super(comparator);
    }

    public RBTree() {
        this(null);
    }

    @Override
    public void afterPut(Node<E> node) {
        Node<E> parent = node.parent;
        if (parent == null) {
            black(node);
            return;
        }
        if (isBlack(parent)) {
            return;
        }
        Node<E> uncle = parent.sibling();
        Node<E> grand = red(parent.parent);
        if (isRed(uncle)) {
            black(uncle);
            black(parent);
            afterPut(grand);
            return;
        }
        if (parent.isLeftChild()) {
            if (node.isRightChild()) {
                black(node);
                rotateLeft(parent);
            } else {
                black(parent);
            }
            rotateRight(grand);
        } else {
            if (node.isLeftChild()) {
                black(node);
                rotateRight(parent);
            } else {
                black(parent);
            }
            rotateLeft(grand);
        }
    }


    @Override
    public void afterRemove(Node<E> node) {
        Node<E> parent = node.parent;
        if (parent == null) {
            return;
        }
        if (isRed(node)) {
            black(node);
            return;
        }
        boolean left = parent.left == null || node.isLeftChild();
        Node<E> sibling = left ? parent.right : parent.left;
        if (left) {
            if (isRed(sibling)) {
                black(sibling);
                red(parent);
                rotateLeft(parent);
                sibling = parent.right;
            }
            if (isBlack(sibling.right) && isBlack(sibling.left)) {
                boolean parentColor = isBlack(parent);
                red(sibling);
                black(parent);
                if (parentColor) {
                    afterRemove(parent);
                }
            } else {
                if (isBlack(sibling.right)) {
                    rotateRight(sibling);
                    sibling = parent.right;
                }
                colorOf(sibling, color(parent));
                black(parent);
                black(sibling.right);
                rotateLeft(parent);
            }
        } else {
            if (isRed(sibling)) {
                black(sibling);
                red(parent);
                rotateRight(parent);
                sibling = parent.left;
            }
            if (isBlack(sibling.right) && isBlack(sibling.left)) {
                boolean parentColor = isBlack(parent);
                red(sibling);
                black(parent);
                if (parentColor) {
                    afterRemove(parent);
                }
            } else {
                if (isBlack(sibling.left)) {
                    rotateLeft(sibling);
                    sibling = parent.left;
                }
                colorOf(sibling, color(parent));
                black(parent);
                black(sibling.left);
                rotateRight(parent);
            }
        }
    }

    @Override
    public Node<E> createNode(E element, Node<E> parent) {
        return new RBNode<>(element, parent);
    }

    public Node<E> colorOf(Node<E> node, boolean color) {
        if (node == null) {
            return null;
        }
        ((RBNode<E>) node).color = color;
        return node;
    }

    public boolean color(Node<E> node) {
        if (node == null) {
            return BLACK;
        }
        return ((RBNode<E>) node).color;
    }

    public Node<E> red(Node<E> node) {
        return colorOf(node, RED);
    }

    public Node<E> black(Node<E> node) {
        return colorOf(node, BLACK);
    }

    public boolean isRed(Node<E> node) {
        return color(node) == RED;
    }

    public boolean isBlack(Node<E> node) {
        return color(node) == BLACK;
    }

    public static class RBNode<E> extends Node<E>{

        boolean color = RED;

        public RBNode(E element, Node<E> parent) {
            super(element, parent);
        }
    }
}
